package registry

import (
	"context"
	"fmt"
	clientv3 "go.etcd.io/etcd/client/v3"
	"strings"
	"time"
)

const EtcdRegistryContextKey = "etcd-registry"

type EtcdRegistry struct {
	Client *clientv3.Client
}

func NewEtcdRegistry(endpoints []string) (*EtcdRegistry, error) {
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: 5 * time.Second})
	if err != nil {
		return nil, err
	}
	return &EtcdRegistry{
		Client: cli,
	}, nil
}

func (r *EtcdRegistry) RegisterService(serviceName string, serviceAddress string) error {
	leaseGrantResp, err := r.Client.Lease.Grant(context.TODO(), 10)
	if err != nil {
		return err
	}
	key := fmt.Sprintf("services/%s/%s", serviceName, serviceAddress)
	_, err = r.Client.Put(context.TODO(), key, "alive", clientv3.WithLease(leaseGrantResp.ID))
	return err
}

func (r *EtcdRegistry) DiscoverServiceInstances(serviceName string) ([]string, error) {
	resp, err := r.Client.Get(context.TODO(), fmt.Sprintf("service/%s/", serviceName), clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	var services []string
	for _, ev := range resp.Kvs {
		parts := strings.Split(string(ev.Key), "/")
		if len(parts) >= 3 {
			services = append(services, parts[2])
		}
	}
	return services, nil
}
